import numpy as np
import tensorflow as tf

# 加载词表
word2idx = {w: i for i, w in enumerate(open('./vocab/word.txt', encoding='utf-8').read().splitlines())}
vocab_size = len(word2idx) + 1
embedding = np.load('./vocab/word.npy')

# 定义模型（要和训练时一致）
class Model(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.embedding = tf.Variable(embedding, dtype=tf.float32, trainable=False)
        self.rnn1 = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(200, return_sequences=True))
        self.rnn2 = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(200, return_sequences=True))
        self.rnn3 = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(200, return_sequences=False))
        self.fc = tf.keras.layers.Dense(400, activation=tf.nn.elu)
        self.out_linear = tf.keras.layers.Dense(2)

    def call(self, inputs, training=False):
        if inputs.dtype != tf.int32:
            inputs = tf.cast(inputs, tf.int32)
        x = tf.nn.embedding_lookup(self.embedding, inputs)
        x = self.rnn1(x)
        x = self.rnn2(x)
        x = self.rnn3(x)
        x = self.fc(x)
        x = self.out_linear(x)
        return x

# 载入模型
model = Model()
model.build(input_shape=(None, None))
model.load_weights('./models/best_model/variables/variables')

# 预测函数
def predict(sentence):
    # 预处理
    words = sentence.lower().split()
    ids = [word2idx.get(w, len(word2idx)) for w in words]  # 如果词表里没有，就设成未知
    max_len = 1000
    if len(ids) >= max_len:
        ids = ids[:max_len]
    else:
        ids += [0] * (max_len - len(ids))
    ids = np.array([ids], dtype=np.int32)  # batch size = 1

    # 推理
    logits = model(ids, training=False)
    pred = tf.argmax(logits, axis=-1).numpy()[0]
    return pred

# 示例
if __name__ == "__main__":
    while True:
        sentence = input("请输入英文句子（输入q退出）: ")
        if sentence.lower() == 'q':
            break
        label = predict(sentence)
        print(f"预测标签: {label} ({'正面' if label == 1 else '负面'})")
